import sys
import glob
import os

DataDirPath = sys.argv[1]
workdir = sys.argv[2]   #not sample workdir.run cellrange will create new sampledir
CITEBar = sys.argv[3]  #if no,fill "F"
HashTag = sys.argv[4]  #if no,fill "F"

#finish in 2021.11.8
#fastq: SP fastqs can't obtain "CDNA"!

#CITESeq or NO CITESeq both ok!

#base function
def getDirPath(DirPath):
    tempDir = glob.glob(DirPath)
    tempDir.sort(reverse=False)
    return tempDir

#base function end.

#main
#cellranger = "/disk/dongxj/software/cellranger-3.0.2/cellranger"
cellranger = "/data1/home/zhoulab/software/3.cellranger6/cellranger-6.1.1/cellranger"
RefCDNA = "/data1/home/zhoulab/database/cellranger/refdata-cellranger-GRCh38-3.0.0"
RefXCR = "/data1/home/zhoulab/database/cellranger/refdata-cellranger-vdj-GRCh38-alts-ensembl-2.0.0"

head = "#BSUB -n 16\n#BSUB -o ddd.our\n#BSUB -e ddd.err\n#BSUB -R span[hosts=1]\n#BSUB -q smp\n#######BSUB -m s001\n"

flag = 0
hflag = 0
CDNADirPath = "%s/*CDNA*" % (DataDirPath)
SPDirPath = "%s/*SP*" % (DataDirPath)
CDNADir = getDirPath(CDNADirPath)
SPDir = getDirPath(SPDirPath)

TCRDirPath = "%s/*TCR*" % (DataDirPath)
TCRDir = getDirPath(TCRDirPath)
BCRDirPath = "%s/*BCR*" % (DataDirPath)
BCRDir = getDirPath(BCRDirPath)

if len(CDNADir) > 0 and len(SPDir) == 0:
    flag = 1   #only CDNA
    print("find only CDNA")
elif len(CDNADir) > 0 and len(SPDir) > 0:
    flag = 2   #CITESEQ
    print("find CDNA+CITESEQ")
elif len(TCRDir) > 0 and len(BCRDir) > 0:
    flag = 5
    print("find TCR+BCR")
elif len(TCRDir) > 0:
    flag = 3
    print("find TCR")
elif len(BCRDir) > 0:
    flag = 4
    print("find BCR")
else:
    print("please check fastq name contain CDNA,TCR,BCR,SP!!!")

if os.path.isfile(HashTag):
    hflag = 1
    hashline = []
    f = open(HashTag)
    line = f.readline()
    while line:
        hashline.append(line)
        line = f.readline()
    f.close()
#
print("flag:",flag,"hashtag flag:",hflag)

if flag == 2:#CITESEQ
   sampleName = os.path.basename(CDNADir[0])
   sampleName = sampleName.split("_")[0]

   SpName = os.path.basename(SPDir[0])
   SpName = SpName.split("_")[0]

   print("CDNA name:",sampleName,"SP name:",SpName)

   #create library File
   libFile = "%s/CITE_%s_library.csv" % (workdir,sampleName)
   OutLines = "fastqs,sample,library_type\n"
   for i in CDNADir:
       stri = "%s/,%s,Gene Expression\n" % (i,sampleName)
       OutLines += stri

   for i in SPDir:
       stri = "%s/,%s,Antibody Capture\n" % (i,SpName)
       OutLines += stri

   w = open(libFile,mode="w+")
   w.write(OutLines)
   w.close()
   #create library File end.

   #create barcode_ref file
   barFile = "%s/CITE_%s_barcoderef.csv" % (workdir,sampleName)
   OutLines = "id,name,read,pattern,sequence,feature_type\n"
   
   #citeseq
   if CITEBar != "F":
       CITEBarTemp = CITEBar.split(",")
       c = 1
       for i in CITEBarTemp:
           CITEBarName = "CITE-%s-%s" % (sampleName,c)
           stri = "%s,%s,R2,NNNNNNNNNN(BC),%s,Antibody Capture\n" % (CITEBarName,CITEBarName,i)
           OutLines += stri
           c += 1
       print("find citeseq:",c)
   else:
       print("no find citeseq!")

   #protein hash tag
   if hflag == 1:
       c = 0
       for i in hashline:
           OutLines += i
           c += 1
       print("find protein hash tag:",c)
   else:
       print("no find protein hash tag!!")
   
   #write
   w = open(barFile,mode="w+")
   w.write(OutLines) 
   w.close()
 
   #CDNA run sh
   cmdsh = "%s/run_%s-CDNA.sh" % (workdir,sampleName) 
   cmd = "%s count --id=%s --transcriptome=%s --libraries=%s --feature-ref=%s" % (cellranger,sampleName,RefCDNA,libFile,barFile)
   w = open(cmdsh,mode="w+")
   w.write(head)
   w.write(cmd)
   w.close()
   #CDNA run sh end.


elif flag == 1: #no CITESeq no hashtag
   print("no CITESeq and no protein hash tag!")

   sampleName = os.path.basename(CDNADir[0])
   sampleName = sampleName.split("_")[0]

   fqstri = ""
   for i in CDNADir:
       stri = "%s," % (i)
       fqstri += stri
   fqstri = fqstri.strip(",")

   cmd = "%s count --id=%s --transcriptome=%s --fastqs=%s --sample=%s --project=%s" % (cellranger,sampleName,RefCDNA,fqstri,sampleName,sampleName)

   #CDNA run sh
   cmdsh = "%s/run_%s-CDNA.sh" % (workdir,sampleName)
   w = open(cmdsh,mode="w+")
   w.write(head)
   w.write(cmd)
   #print(cmd)
   w.close()
   #CDNA run sh end.

elif flag == 3 or flag == 4:
   print("TCR")

   if flag == 3:
       TBCRDir = TCRDir
   else:
       TBCRDir = BCRDir

   sampleName = os.path.basename(TBCRDir[0])
   sampleName = sampleName.split("_")[0]

   fqstri = ""
   for i in TBCRDir:
       stri = "%s," % (i)
       fqstri += stri
   fqstri = fqstri.strip(",")

   cmd = "%s vdj --id=%s --reference=%s --fastqs=%s --sample=%s --project=%s" % (cellranger,sampleName,RefXCR,fqstri,sampleName,sampleName)

   #TCR/BCR run sh
   if flag == 3:
       cmdsh = "%s/run_%s-TCR.sh" % (workdir,sampleName)
   else:
       cmdsh = "%s/run_%s-BCR.sh" % (workdir,sampleName)
   w = open(cmdsh,mode="w+")
   w.write(head)
   w.write(cmd)
   w.close()
   #TCR/BCR run sh end.

elif flag == 5: #TCR and BCR

   sampleName = os.path.basename(TCRDir[0])
   sampleName = sampleName.split("_")[0]

   fqstri = ""
   for i in TCRDir:
       stri = "%s," % (i)
       fqstri += stri
   fqstri = fqstri.strip(",")

   cmd = "%s vdj --id=%s --reference=%s --fastqs=%s --sample=%s --project=%s" % (cellranger,sampleName,RefXCR,fqstri,sampleName,sampleName)

   cmdsh = "%s/run_%s-TCR.sh" % (workdir,sampleName)
   w = open(cmdsh,mode="w+")
   w.write(head)
   w.write(cmd)
   w.close()

   #BCR
   sampleName = os.path.basename(BCRDir[0])
   sampleName = sampleName.split("_")[0]

   fqstri = ""
   for i in BCRDir:
       stri = "%s," % (i)
       fqstri += stri
   fqstri = fqstri.strip(",")

   cmd = "%s vdj --id=%s --reference=%s --fastqs=%s --sample=%s --project=%s" % (cellranger,sampleName,RefXCR,fqstri,sampleName,sampleName)

   cmdsh = "%s/run_%s-BCR.sh" % (workdir,sampleName)
   w = open(cmdsh,mode="w+")
   w.write(head)
   w.write(cmd)
   w.close()


#test run
#python CDNA-ZDFX.py /data1/home/zhoulab/testfq /data1/home/zhoulab/workdir AAGTATCGTTTCGCA
